#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
WebSocket Server for Smart Home Butler ASR System
智能管家语音识别系统的WebSocket服务器
"""

import asyncio
import websockets
import json
import base64
import logging
import os
from queue import Queue
import threading
from funasr.utils.postprocess_utils import rich_transcription_postprocess
from .optimized_realtime_asr import OptimizedRealTimeASR, ASRProcessor, ContextManager, SemanticSegmenter

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class ASROutputHandler:
    """
    ASR输出处理器
    处理ASR系统的输出并发送到WebSocket客户端
    """
    
    def __init__(self, websocket):
        self.websocket = websocket
        self.last_text = ""
    
    def handle_result(self, text):
        """
        处理识别结果
        :param text: 识别的文本
        """
        if text and text != self.last_text:
            self.last_text = text
            # 创建异步任务发送结果
            asyncio.create_task(self.send_result(text))
    
    async def send_result(self, text):
        """
        发送结果到客户端
        :param text: 识别的文本
        """
        try:
            response = {
                "type": "result",
                "text": text,
                "is_complete": True  # 简化处理，始终标记为完整句子
            }
            await self.websocket.send(json.dumps(response))
        except Exception as e:
            logger.error(f"发送结果到客户端时出错: {str(e)}")


class ASRWebSocketSession:
    """
    ASR WebSocket会话
    管理单个客户端连接的ASR处理流程
    """
    
    def __init__(self, websocket, model_dir="iic/SenseVoiceSmall", model_cache_dir=None):
        self.websocket = websocket
        self.model_dir = model_dir
        self.model_cache_dir = model_cache_dir
        self.asr_system = None
        self.output_handler = ASROutputHandler(websocket)
        self.audio_buffer = bytearray()
        self.is_processing = False
        
    def initialize_asr(self):
        """
        初始化ASR系统
        """
        try:
            # 创建ASR系统实例
            self.asr_system = OptimizedRealTimeASR(model_dir=self.model_dir, model_cache_dir=self.model_cache_dir)
            
            # 替换输出处理函数
            original_process_audio = self.asr_system.process_audio
            
            def custom_process_audio():
                # 这里我们不需要完全替换原有的处理逻辑
                # 只是在适当的时候捕获输出
                original_process_audio()
                
            self.asr_system.process_audio = custom_process_audio
            
            logger.info("ASR系统初始化成功")
            return True
        except Exception as e:
            logger.error(f"ASR系统初始化失败: {str(e)}")
            return False
    
    async def process_audio_chunk(self, audio_data):
        """
        处理音频块
        :param audio_data: 音频数据
        """
        try:
            # 发送处理中状态
            await self.websocket.send(json.dumps({
                "type": "status",
                "message": "processing"
            }))
            
            # 对于WebSocket服务，我们直接调用ASR处理器
            # 创建必要的组件
            context_manager = ContextManager()
            semantic_segmenter = SemanticSegmenter()
            
            # 导入必要的模块
            from funasr import AutoModel
            import torch
            import numpy as np
            
            # 设置模型缓存目录
            if self.model_cache_dir:
                os.environ['MODELSCOPE_CACHE'] = self.model_cache_dir
            
            # 初始化模型，禁用更新检查
            model = AutoModel(
                model=self.model_dir,
                trust_remote_code=True,
                vad_model="iic/speech_fsmn_vad_zh-cn-16k-common-pytorch",
                vad_kwargs={"max_single_segment_time": 30000},
                device="cuda:0" if torch.cuda.is_available() else "cpu",
                disable_pbar=True,
                disable_update=True  # 禁用模型更新检查
            )
            
            # 初始化ASR处理器
            asr_processor = ASRProcessor(model, context_manager, semantic_segmenter)
            
            # 处理音频数据 - 注意这里需要传入列表格式
            result, is_complete = asr_processor.recognize_speech([audio_data])
            
            if result is not None:
                # 发送识别结果
                response = {
                    "type": "result",
                    "text": result,
                    "is_complete": is_complete
                }
                await self.websocket.send(json.dumps(response))
            else:
                # 发送空结果状态
                await self.websocket.send(json.dumps({
                    "type": "status",
                    "message": "no_speech_detected"
                }))
                
        except Exception as e:
            logger.error(f"处理音频时出错: {str(e)}")
            await self.websocket.send(json.dumps({
                "type": "error",
                "message": f"处理音频时出错: {str(e)}"
            }))


class ASRWebSocketServer:
    """
    ASR WebSocket服务器类
    提供基于WebSocket的语音识别服务
    """

    def __init__(self, model_dir="iic/SenseVoiceSmall", model_cache_dir=None, host="localhost", port=8765):
        """
        初始化WebSocket服务器
        :param model_dir: ASR模型路径
        :param model_cache_dir: 模型缓存目录路径
        :param host: 服务器主机地址
        :param port: 服务器端口
        """
        self.host = host
        self.port = port
        self.model_dir = model_dir
        self.model_cache_dir = model_cache_dir
        # 存储客户端会话
        self.sessions = {}
        logger.info(f"ASR WebSocket服务器初始化完成，将监听 {host}:{port}")

    async def register_client(self, websocket):
        """
        注册客户端连接
        :param websocket: WebSocket连接对象
        """
        session = ASRWebSocketSession(websocket, self.model_dir, self.model_cache_dir)
        self.sessions[websocket] = session
        logger.info(f"新客户端连接: {websocket.remote_address}, 当前客户端数量: {len(self.sessions)}")
        return session

    async def unregister_client(self, websocket):
        """
        注销客户端连接
        :param websocket: WebSocket连接对象
        """
        if websocket in self.sessions:
            del self.sessions[websocket]
        logger.info(f"客户端断开连接: {websocket.remote_address}, 剩余客户端数量: {len(self.sessions)}")

    async def handle_audio_data(self, session, audio_data):
        """
        处理音频数据
        :param session: ASR会话对象
        :param audio_data: 音频数据（base64编码）
        """
        try:
            # 解码音频数据
            audio_bytes = base64.b64decode(audio_data)
            
            # 处理音频块
            await session.process_audio_chunk(audio_bytes)
            
        except Exception as e:
            logger.error(f"处理音频数据时出错: {str(e)}")
            await session.websocket.send(json.dumps({
                "type": "error",
                "message": f"处理音频数据时出错: {str(e)}"
            }))

    async def handle_message(self, session, message):
        """
        处理来自客户端的消息
        :param session: ASR会话对象
        :param message: 接收到的消息
        """
        try:
            data = json.loads(message)
            msg_type = data.get("type")
            
            if msg_type == "audio":
                # 处理音频数据
                await self.handle_audio_data(session, data.get("data", ""))
            elif msg_type == "reset_context":
                # 重置上下文（当前实现中每次请求都是独立的）
                await session.websocket.send(json.dumps({
                    "type": "status",
                    "message": "context_reset"
                }))
            else:
                # 未知消息类型
                await session.websocket.send(json.dumps({
                    "type": "error",
                    "message": f"未知消息类型: {msg_type}"
                }))
                
        except json.JSONDecodeError:
            await session.websocket.send(json.dumps({
                "type": "error",
                "message": "无效的JSON格式"
            }))
        except Exception as e:
            logger.error(f"处理消息时出错: {str(e)}")
            await session.websocket.send(json.dumps({
                "type": "error",
                "message": f"处理消息时出错: {str(e)}"
            }))

    async def client_handler(self, websocket, path=None):
        """
        客户端处理器
        :param websocket: WebSocket连接对象
        :param path: 请求路径
        """
        # 注册客户端
        session = await self.register_client(websocket)
        
        try:
            # 发送欢迎消息
            welcome_msg = {
                "type": "welcome",
                "message": "欢迎使用智能管家语音识别服务"
            }
            await websocket.send(json.dumps(welcome_msg))
            
            # 监听客户端消息
            async for message in websocket:
                await self.handle_message(session, message)
                
        except websockets.exceptions.ConnectionClosed:
            logger.info("客户端连接已关闭")
        except Exception as e:
            logger.error(f"客户端处理过程中出错: {str(e)}")
        finally:
            # 注销客户端
            await self.unregister_client(websocket)

    def run(self):
        """
        启动WebSocket服务器
        """
        logger.info(f"启动WebSocket服务器 {self.host}:{self.port}")
        
        # 使用asyncio.run来处理事件循环（推荐方法）
        async def start_server():
            server = await websockets.serve(self.client_handler, self.host, self.port)
            logger.info(f"WebSocket服务器已启动并监听 {self.host}:{self.port}")
            await server.wait_closed()
        
        try:
            asyncio.run(start_server())
        except KeyboardInterrupt:
            logger.info("服务器已被用户中断")
        except Exception as e:
            logger.error(f"服务器运行出错: {str(e)}")
            raise


# 服务器启动函数
def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description="智能管家语音识别WebSocket服务器")
    parser.add_argument("--host", default="localhost", help="服务器主机地址 (默认: localhost)")
    parser.add_argument("--port", type=int, default=8765, help="服务器端口 (默认: 8765)")
    parser.add_argument("--model-dir", default="iic/SenseVoiceSmall", help="ASR模型路径")
    parser.add_argument("--model-cache-dir", default=None, help="模型缓存目录路径")
    
    args = parser.parse_args()
    
    # 设置MODELSCOPE_CACHE环境变量，确保模型缓存到正确的目录
    if args.model_cache_dir:
        # 注意：ModelScope会在指定目录下创建models子目录，所以我们直接使用传入的目录
        os.environ['MODELSCOPE_CACHE'] = args.model_cache_dir
    
    try:
        server = ASRWebSocketServer(
            model_dir=args.model_dir,
            model_cache_dir=args.model_cache_dir,
            host=args.host,
            port=args.port
        )
        server.run()
    except KeyboardInterrupt:
        logger.info("服务器已停止")


if __name__ == "__main__":
    main()